# Simple script to go through a load of projection data (in the format
# used by the PRM as input) and collate all years for a particular region into
# a single file (one line for each year)

import sys
import os.path
import re

__author__="nick"
__date__ ="$Oct 27, 2011 10:18:15 AM$"

usage = """Usage: python prm_input_parser <input files dir> <region name>"""

def sort_files(filelist):
    """Sorts prm input data files by year. Each file should have a
    numeric year in the file name somewhere - this number is used to sort.
    The function assumes that there are no missing years."""

    # Regex to extrat the year from the filename
    regex = re.compile(r".*?(\d\d*).*")

    # First find the start year
    min_year = 9999
    for file in filelist:
        m = regex.match(file)
        assert m != None, "No match on "+file
        yr = int(m.group(1))
        if yr < min_year:
            min_year = yr

    sorted_files = [None for i in range(len(filelist))]

    # Iterate again, putting files in the correct position
    for file in filelist:
        m = regex.match(file)
        assert m != None, "No match on "+file
        yr = int(m.group(1))
        sorted_files[yr-min_year] = file

    return sorted_files


if __name__ == "__main__":

    #print "Arguments",sys.argv

    # Check command-line arguments, should be directory and a number of regions
    if (len(sys.argv)!=3):
        print usage
        sys.exit()

    dir = sys.argv[1]

    if not os.path.isdir(dir):
        print "The first argument must be a directory, not",dir
        print usage
        sys.exit()

    region = sys.argv[2]
    if len(region)!=4:
        print "Second argument should be a four-digit region name, not",region
        print usage
        sys.exit()



    print "Will look for files to collate in the directory",dir

    # List files by year (ascending)

    allfiles = os.listdir(dir)
    sortedfiles = sort_files(allfiles)

    outfilename = "collated_"+region+".csv"
    outfile = open(outfilename, "w")

    firstfile = True # Needed for writing file header
    for infile in sortedfiles:

        f = open(dir+infile, "r")
        foundregion = False
        for line in f:

            if (firstfile): # Need to write the header once
                outfile.write(line)
                firstfile = False # To show that we've read the first file

            # Look for the required region
            if line[0:4] == region:
                foundregion = True
                outfile.write(line)
                #print "Found relevant entry: ",line


        f.close()

        if not foundregion:
            print "Warning, region",region,"was not found in any of the filees."
        
        

    print "Have written a collated file: ",outfile.name
    outfile.close()
